import ctypes
import ctypes.wintypes as wintypes
import struct
import sys
import os
import subprocess
import tempfile

# Windows API常量
ACPI_SIGNATURE = 0x41435049  # 'ACPI'
DMAR_SIGNATURE = 0x52414D44  # 'DMAR'

# 加载kernel32.dll
kernel32 = ctypes.windll.kernel32

class EFI_ACPI_DESCRIPTION_HEADER:
    """EFI_ACPI_DESCRIPTION_HEADER 结构体解析"""
    def __init__(self, data):
        if len(data) < 36:
            raise ValueError("数据长度不足以解析ACPI描述头")
        
        # 解析ACPI标准头部 (36字节)
        header_data = struct.unpack('<4sLBB6sQLLL', data[:36])
        
        self.Signature = header_data[0].decode('ascii', errors='ignore')
        self.Length = header_data[1]
        self.Revision = header_data[2]
        self.Checksum = header_data[3]
        self.OemId = header_data[4].decode('ascii', errors='ignore').rstrip('\x00')
        self.OemTableId = struct.pack('<Q', header_data[5])[:8].decode('ascii', errors='ignore').rstrip('\x00')
        self.OemRevision = header_data[6]
        self.CreatorId = struct.pack('<L', header_data[7]).decode('ascii', errors='ignore')
        self.CreatorRevision = header_data[8]
    
    def print_info(self):
        """打印ACPI描述头信息"""
        print(f"📋 EFI_ACPI_DESCRIPTION_HEADER:")
        print(f"    签名 (Signature): '{self.Signature}'")
        print(f"    长度 (Length): {self.Length} 字节")
        print(f"    修订版 (Revision): {self.Revision}")
        print(f"    校验和 (Checksum): 0x{self.Checksum:02X}")
        print(f"    OEM ID: '{self.OemId}'")
        print(f"    OEM Table ID: '{self.OemTableId}'")
        print(f"    OEM修订版 (OemRevision): {self.OemRevision}")
        print(f"    创建者ID (CreatorId): '{self.CreatorId}'")
        print(f"    创建者修订版 (CreatorRevision): 0x{self.CreatorRevision:08X}")

class EFI_ACPI_DMAR_HEADER:
    """EFI_ACPI_DMAR_HEADER 结构体解析"""
    def __init__(self, data):
        if len(data) < 48:
            raise ValueError("数据长度不足以解析DMAR头")
        
        # 首先解析ACPI描述头
        self.Header = EFI_ACPI_DESCRIPTION_HEADER(data)
        
        # 解析DMAR特定字段 (从偏移36开始，12字节)
        dmar_fields = struct.unpack('<BB10s', data[36:48])
        
        self.HostAddressWidth = dmar_fields[0]
        self.Flags = dmar_fields[1]
        self.Reserved = dmar_fields[2]
    
    def get_flags_description(self):
        """获取标志位的描述"""
        flags_desc = []
        
        if self.Flags & 0x01:  # BIT0
            flags_desc.append("INTR_REMAP (支持中断重映射)")
        if self.Flags & 0x02:  # BIT1
            flags_desc.append("X2APIC_OPT_OUT (请求退出X2APIC模式)")
        if self.Flags & 0x04:  # BIT2
            flags_desc.append("DMA_CTRL_PLATFORM_OPT_IN_FLAG (平台DMA控制)")
        
        return flags_desc
    
    def print_info(self):
        """打印DMAR头信息"""
        # 打印ACPI描述头
        self.Header.print_info()
        
        print(f"\n🔧 EFI_ACPI_DMAR_HEADER 特定字段:")
        print(f"    主机地址宽度 (HostAddressWidth): {self.HostAddressWidth} 位")
        print(f"        -> 实际物理地址宽度: {self.HostAddressWidth + 1} 位")
        print(f"        -> 最大可寻址内存: {2**(self.HostAddressWidth + 1) // (1024**3)} GB")
        
        print(f"    标志 (Flags): 0x{self.Flags:02X}")
        flags_desc = self.get_flags_description()
        if flags_desc:
            for desc in flags_desc:
                print(f"        - {desc}")
        else:
            print(f"        - 无特殊标志")
        
        print(f"    保留字段 (Reserved): {self.Reserved.hex().upper()}")

def parse_dmar_remapping_structures(data, offset=48):
    """解析DMAR重映射结构"""
    print(f"\n🏗️  解析DMAR重映射结构 (从偏移 0x{offset:04X} 开始):")
    
    structures = []
    current_offset = offset
    structure_count = 0
    
    while current_offset < len(data):
        if current_offset + 4 > len(data):
            break
        
        # 读取结构头部
        struct_type = struct.unpack('<H', data[current_offset:current_offset+2])[0]
        struct_length = struct.unpack('<H', data[current_offset+2:current_offset+4])[0]
        
        if struct_length == 0:
            print(f"    ⚠️  警告: 结构长度为0，停止解析")
            break
        
        if current_offset + struct_length > len(data):
            print(f"    ⚠️  警告: 结构长度超出数据范围，停止解析")
            break
        
        structure_count += 1
        print(f"\n  📦 结构 {structure_count} (偏移 0x{current_offset:04X} - 0x{current_offset + struct_length - 1:04X}):")
        print(f"      类型 (Type): 0x{struct_type:04X}", end="")
        
        # 解析具体的结构类型
        if struct_type == 0x0000:  # EFI_ACPI_DMAR_TYPE_DRHD
            print(" - DRHD (DMA重映射硬件单元定义)")
            parse_drhd_structure(data, current_offset, struct_length)
        elif struct_type == 0x0001:  # EFI_ACPI_DMAR_TYPE_RMRR
            print(" - RMRR (保留内存区域报告)")
            parse_rmrr_structure(data, current_offset, struct_length)
        elif struct_type == 0x0002:  # EFI_ACPI_DMAR_TYPE_ATSR
            print(" - ATSR (ATS根端口报告)")
            parse_atsr_structure(data, current_offset, struct_length)
        elif struct_type == 0x0003:  # EFI_ACPI_DMAR_TYPE_RHSA
            print(" - RHSA (重映射硬件静态关联)")
            parse_rhsa_structure(data, current_offset, struct_length)
        elif struct_type == 0x0004:  # EFI_ACPI_DMAR_TYPE_ANDD
            print(" - ANDD (ACPI命名空间设备声明)")
            parse_andd_structure(data, current_offset, struct_length)
        elif struct_type == 0x0005:  # EFI_ACPI_DMAR_TYPE_SATC
            print(" - SATC (SoC集成地址转换缓存)")
            parse_satc_structure(data, current_offset, struct_length)
        elif struct_type == 0x0006:  # EFI_ACPI_DMAR_TYPE_SIDP
            print(" - SIDP (SoC集成设备属性)")
            parse_sidp_structure(data, current_offset, struct_length)
        else:
            print(f" - 未知类型")
        
        print(f"      长度 (Length): {struct_length} 字节")
        
        structures.append({
            'type': struct_type,
            'length': struct_length,
            'offset': current_offset,
            'data': data[current_offset:current_offset + struct_length]
        })
        
        current_offset += struct_length
    
    print(f"\n📊 总共解析了 {structure_count} 个重映射结构")
    return structures

def parse_drhd_structure(data, offset, length):
    """解析DRHD结构 (EFI_ACPI_DMAR_DRHD_HEADER)"""
    if length < 16:
        print(f"      ⚠️  DRHD结构长度不足")
        return
    
    # 解析DRHD字段
    drhd_data = struct.unpack('<HHBBHQ', data[offset:offset+16])
    
    flags = drhd_data[2]
    size = drhd_data[3]
    segment_number = drhd_data[4]
    register_base_address = drhd_data[5]
    
    print(f"      标志 (Flags): 0x{flags:02X}")
    if flags & 0x01:
        print(f"          - INCLUDE_PCI_ALL: 包含段内所有PCI设备")
    
    print(f"      大小 (Size): {size} (寄存器集大小: {2**size * 4} KB)")
    print(f"      段号 (SegmentNumber): {segment_number}")
    print(f"      寄存器基址 (RegisterBaseAddress): 0x{register_base_address:016X}")
    
    # 解析设备范围条目 (如果有)
    if length > 16:
        print(f"      设备范围条目:")
        parse_device_scope_entries(data, offset + 16, length - 16)

def parse_rmrr_structure(data, offset, length):
    """解析RMRR结构 (EFI_ACPI_DMAR_RMRR_HEADER)"""
    if length < 24:
        print(f"      ⚠️  RMRR结构长度不足")
        return
    
    # 解析RMRR字段
    rmrr_data = struct.unpack('<HHHHQQ', data[offset:offset+24])
    
    segment_number = rmrr_data[3]
    base_address = rmrr_data[4]
    limit_address = rmrr_data[5]
    
    print(f"      段号 (SegmentNumber): {segment_number}")
    print(f"      基址 (BaseAddress): 0x{base_address:016X}")
    print(f"      限制地址 (LimitAddress): 0x{limit_address:016X}")
    print(f"      区域大小: {(limit_address - base_address + 1) // 1024} KB")
    
    # 解析设备范围条目 (如果有)
    if length > 24:
        print(f"      设备范围条目:")
        parse_device_scope_entries(data, offset + 24, length - 24)

def parse_atsr_structure(data, offset, length):
    """解析ATSR结构 (EFI_ACPI_DMAR_ATSR_HEADER)"""
    if length < 8:
        print(f"      ⚠️  ATSR结构长度不足")
        return
    
    atsr_data = struct.unpack('<HHBBH', data[offset:offset+8])
    
    flags = atsr_data[2]
    segment_number = atsr_data[4]
    
    print(f"      标志 (Flags): 0x{flags:02X}")
    if flags & 0x01:
        print(f"          - ALL_PORTS: 所有根端口支持ATS")
    
    print(f"      段号 (SegmentNumber): {segment_number}")

def parse_rhsa_structure(data, offset, length):
    """解析RHSA结构 (EFI_ACPI_DMAR_RHSA_HEADER)"""
    if length < 20:
        print(f"      ⚠️  RHSA结构长度不足")
        return
    
    rhsa_data = struct.unpack('<HHLLQL', data[offset:offset+20])
    
    register_base_address = rhsa_data[4]
    proximity_domain = rhsa_data[5]
    
    print(f"      寄存器基址 (RegisterBaseAddress): 0x{register_base_address:016X}")
    print(f"      邻近域 (ProximityDomain): {proximity_domain}")

def parse_andd_structure(data, offset, length):
    """解析ANDD结构 (EFI_ACPI_DMAR_ANDD_HEADER)"""
    if length < 8:
        print(f"      ⚠️  ANDD结构长度不足")
        return
    
    andd_data = struct.unpack('<HHBBBB', data[offset:offset+8])
    
    acpi_device_number = andd_data[5]
    
    print(f"      ACPI设备号 (AcpiDeviceNumber): {acpi_device_number}")

def parse_satc_structure(data, offset, length):
    """解析SATC结构 (EFI_ACPI_DMAR_SATC_HEADER)"""
    if length < 8:
        print(f"      ⚠️  SATC结构长度不足")
        return
    
    satc_data = struct.unpack('<HHBBH', data[offset:offset+8])
    
    flags = satc_data[2]
    segment_number = satc_data[4]
    
    print(f"      标志 (Flags): 0x{flags:02X}")
    if flags & 0x01:
        print(f"          - ATC_REQUIRED: ATC为必需")
    
    print(f"      段号 (SegmentNumber): {segment_number}")

def parse_sidp_structure(data, offset, length):
    """解析SIDP结构 (EFI_ACPI_DMAR_SIDP_HEADER)"""
    if length < 8:
        print(f"      ⚠️  SIDP结构长度不足")
        return
    
    sidp_data = struct.unpack('<HHHH', data[offset:offset+8])
    
    segment_number = sidp_data[3]
    
    print(f"      段号 (SegmentNumber): {segment_number}")

def parse_device_scope_entries(data, offset, length):
    """解析设备范围条目 (EFI_ACPI_DMAR_DEVICE_SCOPE_STRUCTURE_HEADER)"""
    current_offset = offset
    entry_count = 0
    
    while current_offset < offset + length:
        if current_offset + 6 > offset + length:
            break
        
        # 解析设备范围头部
        scope_data = struct.unpack('<BBBBBB', data[current_offset:current_offset+6])
        
        scope_type = scope_data[0]
        scope_length = scope_data[1]
        scope_flags = scope_data[2]
        scope_enum_id = scope_data[4]
        start_bus_number = scope_data[5]
        
        entry_count += 1
        
        print(f"        条目 {entry_count}:")
        print(f"          类型 (Type): 0x{scope_type:02X}", end="")
        
        # 解释设备范围类型
        scope_types = {
            0x01: "PCI_ENDPOINT (PCI端点设备)",
            0x02: "PCI_BRIDGE (PCI桥设备)",
            0x03: "IOAPIC (I/O APIC)",
            0x04: "MSI_CAPABLE_HPET (支持MSI的HPET)",
            0x05: "ACPI_NAMESPACE_DEVICE (ACPI命名空间设备)"
        }
        
        if scope_type in scope_types:
            print(f" - {scope_types[scope_type]}")
        else:
            print(f" - 未知类型")
        
        print(f"          长度 (Length): {scope_length}")
        print(f"          标志 (Flags): 0x{scope_flags:02X}")
        print(f"          枚举ID (EnumerationId): {scope_enum_id}")
        print(f"          起始总线号 (StartBusNumber): {start_bus_number}")
        
        # 解析PCI路径 (如果有)
        if scope_length > 6:
            path_length = scope_length - 6
            path_count = path_length // 2
            if path_count > 0:
                print(f"          PCI路径:")
                for i in range(path_count):
                    path_offset = current_offset + 6 + i * 2
                    if path_offset + 2 <= offset + length:
                        device, function = struct.unpack('<BB', data[path_offset:path_offset+2])
                        print(f"            [{i}] 设备: {device}, 功能: {function}")
        
        current_offset += scope_length
        
        if scope_length == 0:
            break

def analyze_minivisor_characteristics(dmar_header):
    """分析MiniVisor特征"""
    print(f"\n🔍 MiniVisor特征分析:")
    
    detected_features = []
    confidence = 0
    
    # 检查OEM Table ID
    if dmar_header.Header.OemTableId == "EDK2":
        detected_features.append("🚨 OEM Table ID为 'EDK2' - 典型的MiniVisor标识")
        confidence += 60
    
    # 检查长度特征
    if dmar_header.Header.Length == 96:
        detected_features.append("⚠️  DMAR表长度为96字节 - 可能是增强版仿真")
        confidence += 30
    elif dmar_header.Header.Length == 80:
        detected_features.append("ℹ️  DMAR表长度为80字节 - 标准长度")
    
    # 检查创建者修订版
    if dmar_header.Header.CreatorRevision == 0x01000013:
        detected_features.append("⚠️  创建者修订版为 0x01000013 - 可能的仿真特征")
        confidence += 20
    
    # 检查主机地址宽度
    if dmar_header.HostAddressWidth == 46:
        detected_features.append("⚠️  主机地址宽度为46位 - MiniVisor典型配置")
        confidence += 15
    
    # 检查标志
    if dmar_header.Flags == 0x05:  # 特定的标志组合
        detected_features.append("⚠️  标志组合可能表示仿真环境")
        confidence += 10
    
    # 输出检测结果
    if detected_features:
        for feature in detected_features:
            print(f"    {feature}")
    else:
        print(f"    ✅ 未检测到明显的MiniVisor特征")
    
    print(f"\n🎯 检测结果:")
    print(f"    置信度: {confidence}%")
    
    if confidence >= 70:
        print(f"    🚨 高风险: 强烈怀疑存在MiniVisor VT-d仿真")
    elif confidence >= 40:
        print(f"    ⚠️  中等风险: 可能存在MiniVisor VT-d仿真")
    elif confidence >= 20:
        print(f"    💛 低风险: 发现可疑特征")
    else:
        print(f"    ✅ 低风险: 未发现明显仿真特征")
    
    return confidence >= 40

def extract_dmar_table():
    """提取DMAR表并保存为文件"""
    try:
        # 获取DMAR表大小
        size = kernel32.GetSystemFirmwareTable(
            ACPI_SIGNATURE,
            DMAR_SIGNATURE,
            None,
            0
        )
        
        if size == 0:
            print("❌ 未找到DMAR表")
            print("可能原因：")
            print("  1. VT-d未启用")
            print("  2. 系统不支持VT-d")
            print("  3. 需要管理员权限")
            return False
        
        print(f"✅ 找到DMAR表，大小: {size} 字节")
        
        # 分配缓冲区
        buffer = (ctypes.c_ubyte * size)()
        
        # 读取DMAR表
        result = kernel32.GetSystemFirmwareTable(
            ACPI_SIGNATURE,
            DMAR_SIGNATURE,
            buffer,
            size
        )
        
        if result == 0:
            print("❌ 读取DMAR表失败")
            return False
        
        # 转换为bytes
        dmar_data = bytes(buffer)
        
        try:
            # 解析DMAR头部结构
            dmar_header = EFI_ACPI_DMAR_HEADER(dmar_data)
            dmar_header.print_info()
            
            # 解析重映射结构
            structures = parse_dmar_remapping_structures(dmar_data)
            
            # 分析MiniVisor特征
            minivisor_detected = analyze_minivisor_characteristics(dmar_header)
            
        except Exception as e:
            print(f"❌ 解析DMAR结构时出错: {e}")
            return False
        
        # 保存到文件
        dmar_file = 'dmar.aml'
        with open(dmar_file, 'wb') as f:
            f.write(dmar_data)
        
        print(f"\n💾 DMAR表已保存到 {dmar_file}")
        
        # 生成完整的十六进制转储 (前640字节或整个文件)
        dump_size = min(640, len(dmar_data))
        print(f"\n📄 十六进制转储 (前{dump_size}字节):")
        print("=" * 80)
        
        for i in range(0, dump_size, 16):
            # 获取当前行的字节
            line_bytes = dmar_data[i:i+16]
            
            # 生成十六进制部分
            hex_part = ' '.join(f'{b:02X}' for b in line_bytes)
            hex_part = hex_part.ljust(48)  # 确保对齐
            
            # 生成ASCII部分
            ascii_part = ''.join(chr(b) if 32 <= b <= 126 else '.' for b in line_bytes)
            
            # 打印格式化的行
            print(f"  {i:04X}: {hex_part} {ascii_part}")
        
        print("=" * 80)
        
        # 如果文件大于640字节，显示统计信息
        if len(dmar_data) > 640:
            print(f"📊 文件总大小: {len(dmar_data)} 字节")
            print(f"📊 已显示: {dump_size} 字节 ({dump_size/len(dmar_data)*100:.1f}%)")
        
        # 保存完整的十六进制转储到文件
        hex_dump_file = 'dmar_hexdump.txt'
        with open(hex_dump_file, 'w', encoding='utf-8') as f:
            f.write(f"DMAR表十六进制转储\n")
            f.write(f"文件大小: {len(dmar_data)} 字节\n")
            f.write("=" * 80 + "\n")
            
            for i in range(0, len(dmar_data), 16):
                line_bytes = dmar_data[i:i+16]
                hex_part = ' '.join(f'{b:02X}' for b in line_bytes)
                hex_part = hex_part.ljust(48)
                ascii_part = ''.join(chr(b) if 32 <= b <= 126 else '.' for b in line_bytes)
                f.write(f"{i:04X}: {hex_part} {ascii_part}\n")
        
        print(f"💾 完整十六进制转储已保存到 {hex_dump_file}")
        
        return True
        
    except Exception as e:
        print(f"❌ 错误: {e}")
        return False

def check_system_requirements():
    """检查系统要求"""
    print("🔍 检查系统要求...")
    
    # 检查管理员权限
    try:
        is_admin = ctypes.windll.shell32.IsUserAnAdmin()
        if is_admin:
            print("✅ 管理员权限: 正常")
        else:
            print("❌ 管理员权限: 需要以管理员身份运行")
            return False
    except:
        print("⚠️  无法检查管理员权限")
    
    # 检查Windows版本
    version = sys.version_info
    print(f"🐍 Python版本: {version.major}.{version.minor}.{version.micro}")
    
    # 检查架构
    architecture = struct.calcsize("P") * 8
    print(f"🏗️  系统架构: {architecture}位")
    
    return True

def main():
    """主函数"""
    print("=" * 80)
    print("🔬 DMAR表结构化分析工具 (完整版)")
    print("   支持完整的EFI_ACPI_DMAR_HEADER和EFI_ACPI_DESCRIPTION_HEADER解析")
    print("=" * 80)
    
    if not check_system_requirements():
        print("\n请以管理员身份重新运行此脚本")
        input("按回车键退出...")
        return
    
    if extract_dmar_table():
        print("\n✅ DMAR表结构化分析完成!")
        print("💡 生成的文件:")
        print("   - dmar.aml (原始ACPI表)")
        print("   - dmar_hexdump.txt (完整十六进制转储)")
    else:
        print("\n❌ DMAR表提取失败")
    
    print("\n" + "=" * 80)
    input("按回车键退出...")

if __name__ == "__main__":
    main()